#include <sstream>
#include <iomanip>

#include "base_util/utils.h"
// #include "common/transforms.h"
#include "common/model_config.h"
#include "det_predictor.h"
#include <opencv2/dnn.hpp>

namespace ai {

#define CYCLE_SIZE 3

static inline float fast_exp(float x) {
  // return exp(x);
  union
  {
    uint32_t i;
    float f;
  } v;
  v.i = (12102203.1616540672 * x + 1064807160.56887296);
  return v.f;
}

HBDetPredictor::HBDetPredictor(HBPackPredictor* model, HBModelManager* manager, LogInfo *lg):
  HBPredictor(model, manager, lg) {
  cur_block = 0;
}


int HBDetPredictor::RunDet(stream::ImageBlob* blob, std::vector<BaseInfo*>& det_infos) {
  int ret = 0;
  // int64_t tmp_time1 = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();

  double t1 = 0;
  if (log_ifo->log_level_4) { t1 = cv::getTickCount(); }
  hbDNNTensor ori_tensor;
  ori_tensor.properties.tensorType = HB_DNN_IMG_TYPE_NV12;
  ori_tensor.properties.tensorLayout = HB_DNN_LAYOUT_NCHW;
  ori_tensor.properties.validShape.dimensionSize[0] = 1;
  ori_tensor.properties.validShape.dimensionSize[1] = 3;
  ori_tensor.properties.validShape.dimensionSize[2] = blob->ori_im_shape[0];
  ori_tensor.properties.validShape.dimensionSize[3] = blob->ori_im_shape[1];
  // Align by 16 bytes
  ori_tensor.properties.alignedShape.numDimensions = 4;
  ori_tensor.properties.alignedShape.dimensionSize[0] = 1;
  ori_tensor.properties.alignedShape.dimensionSize[1] = 3;
  ori_tensor.properties.alignedShape.dimensionSize[2] = blob->ori_im_shape[0];
  ori_tensor.properties.alignedShape.dimensionSize[3] = (blob->ori_im_shape[1]+16-1) / 16 * 16;
  // printf("[liufeng debug] dimensionSize[2]:%d dimensionSize[3]:%d \n",ori_tensor.properties.alignedShape.dimensionSize[2],ori_tensor.properties.alignedShape.dimensionSize[3]);
  int ori_length = ori_tensor.properties.alignedShape.dimensionSize[2] * ori_tensor.properties.alignedShape.dimensionSize[3] * 3 / 2;
  hbSysAllocCachedMem(&ori_tensor.sysMem[0], ori_length);
  memcpy(ori_tensor.sysMem[0].virAddr, blob->yuv.data(), blob->yuv_size);
  hbSysFlushMem(&(ori_tensor.sysMem[0]), HB_SYS_MEM_CACHE_CLEAN);
  hbDNNTaskHandle_t task_handle = nullptr;
  hbDNNResizeCtrlParam ctrl;
  HB_DNN_INITIALIZE_RESIZE_CTRL_PARAM(&ctrl);
  ret = hbDNNResize(&task_handle, &mdl_hb->input_tensors[0], &ori_tensor, nullptr, &ctrl);
  if (log_ifo->log_level_4)spdlog::get("logger")->info("1.2.0.1 hbDNNResize ret: {}", ret);
  hbDNNWaitTaskDone(task_handle, 0);
  hbDNNReleaseTask(task_handle);
  hbSysFreeMem(&(ori_tensor.sysMem[0]));

  if (log_ifo->log_level_4) {
    double tt = ((double)cv::getTickCount() - t1) / cv::getTickFrequency();
    spdlog::get("logger")->info("1.2.0.1 hbDNNResize Run Time: {}, tensorLayout:{}, tensorType:{}", tt, mdl_hb->input_tensors[0].properties.tensorLayout, mdl_hb->input_tensors[0].properties.tensorType);
    spdlog::get("logger")->info("1.2.0.1 validShape.dimensionSize:{}, {}, {}, {}", mdl_hb->input_tensors[0].properties.validShape.dimensionSize[0], mdl_hb->input_tensors[0].properties.validShape.dimensionSize[1], mdl_hb->input_tensors[0].properties.validShape.dimensionSize[2], mdl_hb->input_tensors[0].properties.validShape.dimensionSize[3]);
    spdlog::get("logger")->info("1.2.0.1 alignedShape.dimensionSize:{}, {}, {}, {}\n", mdl_hb->input_tensors[0].properties.alignedShape.dimensionSize[0], mdl_hb->input_tensors[0].properties.alignedShape.dimensionSize[1], mdl_hb->input_tensors[0].properties.alignedShape.dimensionSize[2], mdl_hb->input_tensors[0].properties.alignedShape.dimensionSize[3]);
  }
  if (log_ifo->log_level_5) {
    int32_t yuv_size = mdl_hb->cfg->input_shape[0] * mdl_hb->cfg->input_shape[1] * 3 / 2;
    std::ofstream imageyuv(std::to_string(log_ifo->counter())+".resize.yuv", std::ios::binary);
    if(imageyuv.is_open())
    {
      imageyuv.write((char *)mdl_hb->input_tensors[0].sysMem[0].virAddr, yuv_size);
      imageyuv.close();
    }
  }

  // rec_mut.lock();
  // int64_t tmp_time2 = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();

  mdl_hb->ori_im_shapes = blob->ori_im_shape;
  task_handle = nullptr;
  hbDNNInferCtrlParam infer_ctrl_param;
  HB_DNN_INITIALIZE_INFER_CTRL_PARAM(&infer_ctrl_param);
  hbDNNTensor *output_tensors = mdl_hb->output_tensors.data();
  ret = hbDNNInfer(&task_handle, &output_tensors, mdl_hb->input_tensors.data(), mdl_hb->dnn_handle, &infer_ctrl_param);
  if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.0.2 hbDNNInfer ret: {}", ret);

  hbDNNWaitTaskDone(task_handle, 0);
  for (int i = 0; i < mdl_hb->output_count; i++) {
    hbSysFlushMem(&mdl_hb->output_tensors[i].sysMem[0], HB_SYS_MEM_CACHE_INVALIDATE);
  }
  hbDNNReleaseTask(task_handle);    // 释放task handle
  // rec_mut.unlock();
  if (log_ifo->log_level_4) {
    double tt = ((double)cv::getTickCount() - t1) / cv::getTickFrequency();
    spdlog::get("logger")->info("1.2.0.2 hbDNNInfer Run Time: {}", tt);
  }
  if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.0.2 rec_img");
  // int64_t tmp_time3 = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();

  if (mdl_hb->cfg->arch == "YOLOV5") {
    PostProcess_yolov5(det_infos);
  } 
  else if (mdl_hb->cfg->arch == "YOLOV8") {
    PostProcess_yolov8(blob, det_infos);
  } 
  else {
    spdlog::get("logger")->info("ERROR. HBDetPredictor::RunDet. mdl_hb->cfg->arch error: {} ", mdl_hb->cfg->arch);
  }
  // int64_t tmp_time4 = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
  // std::cout << "RunDet: " << tmp_time2 - tmp_time1 << ", " << tmp_time3 - tmp_time2 << ", " << tmp_time4 - tmp_time3 <<  std::endl;

  cur_block++;
  if (cur_block == CYCLE_SIZE) {cur_block=0;} //最多同时存3个块
  return 0;
}

int HBDetPredictor::PostProcess_yolov5(std::vector<BaseInfo*>& det_infos) {
  int ret = 0;
  // auto& pred = mng->predictors[md_idx];
  double t2 = 0;
  if (log_ifo->log_level_4) { t2 = cv::getTickCount(); }

  std::vector<cv::Rect> boxs;
  std::vector<float> scores;
  std::vector<int> labels;
  for (int i = 0; i < mdl_hb->output_count; i++) {
    auto& t = mdl_hb->output_tensors[i];
    hbSysFlushMem(&(t.sysMem[0]), HB_SYS_MEM_CACHE_INVALIDATE);
    auto& td = t.properties.validShape.dimensionSize;
    std::vector<int> feteature_shape(2, 0);
    ret = get_tensor_hw(&t, &feteature_shape[0], &feteature_shape[1]);
    if (ret != 0) { spdlog::get("logger")->info("Yolo5_detection_parser error"); }
    // spdlog::get("logger")->info("feteature_shape: {}, {}", feteature_shape[0], feteature_shape[1]);

    ParseTensor((float*)t.sysMem[0].virAddr, i, boxs, scores, labels, mdl_hb->ori_im_shapes, feteature_shape, mdl_hb->cfg);
    // spdlog::get("logger")->info("boxs: {}, scores: {}, labels: {}", boxs.size(), scores.size(), labels.size());
  }
  std::vector<int> indices;
  cv::dnn::NMSBoxes(boxs, scores, mdl_hb->cfg->draw_threshold, mdl_hb->cfg->nms_threshold, indices, 100);
  for (auto& idx : indices) {
    DetInfo* di = new DetInfo(scores[idx], labels[idx], boxs[idx], mdl_hb->cfg->label_list[labels[idx]]);
    det_infos.emplace_back(di);
  }

  if (log_ifo->log_level_4) {
    double tt = ((double)cv::getTickCount() - t2) / cv::getTickFrequency();
    spdlog::get("logger")->info("1.2.0.3 NMSBoxes Run Time: {}", tt);
  }
  if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.0.3 det_infos:{}", det_infos.size());
  return 0;
}


float HBDetPredictor::sigmoid(float x)
{
  return 1 / (1 + fast_exp(-x));
}

int HBDetPredictor::GenerateMeshGrid() {
  int ret = 0;
  HeadNum = mdl_hb->cfg->map_size.size();

  for (int index = 0; index < HeadNum; index++) {
    for (int i = 0; i < mdl_hb->cfg->map_size[index][0]; i++) {
      for (int j = 0; j < mdl_hb->cfg->map_size[index][1]; j++) {
        MeshGrid.push_back(float(j + 0.5));
        MeshGrid.push_back(float(i + 0.5));
      }
    }
  }
  return ret;
}
int HBDetPredictor::PostProcess_yolov8(stream::ImageBlob* img_blob, std::vector<BaseInfo*>& det_infos)
{

  int ret = 0;
  if (MeshGrid.empty()) { ret = GenerateMeshGrid(); }

  float w_scale = mdl_hb->cfg->input_shape[1] * 1.0 / mdl_hb->ori_im_shapes[1];
  float h_scale = mdl_hb->cfg->input_shape[0] * 1.0 / mdl_hb->ori_im_shapes[0];
  
  std::vector<cv::Rect> boxs;
  std::vector<float> confidences;
  std::vector<int> labels;

  std::vector<float> RegDfl;
  std::vector<float> RegDeq(16, 0);
  int gridIndex = -2;
  float xmin = 0, ymin = 0, xmax = 0, ymax = 0;
  float cls_val = 0;
  float cls_max = 0;
  int cls_index = 0;
  int class_num = mdl_hb->cfg->label_list.size();
  int quant_zp_cls = 0, quant_zp_reg = 0, quant_zp_msk;
  float quant_scale_cls = 0, quant_scale_reg = 0, quant_scale_msk = 0;
  auto& Strides = mdl_hb->cfg->strides;
  auto& MapSize = mdl_hb->cfg->map_size;
  // printf("[liufeng debug] class_num:%d HeadNum:%d\n",class_num,HeadNum);
  for (int index = 0; index < HeadNum; index++) {
    auto& reg_t = mdl_hb->output_tensors[index * 2 + 0];
    hbSysFlushMem(&(reg_t.sysMem[0]), HB_SYS_MEM_CACHE_INVALIDATE);
    auto& cls_t = mdl_hb->output_tensors[index * 2 + 1];
    hbSysFlushMem(&(cls_t.sysMem[0]), HB_SYS_MEM_CACHE_INVALIDATE);
    float *reg = (float *)reg_t.sysMem[0].virAddr;
    // printf("[liufeng debug] %.4f %.4f %.4f %.4f %.4f\n",reg[0],reg[1],reg[2],reg[3],reg[4]);
    float *cls = (float *)cls_t.sysMem[0].virAddr;
    // printf("[liufeng debug] %.4f %.4f %.4f %.4f %.4f\n",cls[0],cls[1],cls[2],cls[3],cls[4]);

    float sfsum = 0;
    float locval = 0;
    float locvaltemp = 0;

    // printf("[liufeng debug] h:%d w:%d\n",MapSize[index][0],MapSize[index][1]);

    for (int h = 0; h < MapSize[index][0]; h++) {
      for (int w = 0; w < MapSize[index][1]; w++) {
        gridIndex += 2;
        if (1 == class_num) {
          cls_max = sigmoid(cls[0 * MapSize[index][0] * MapSize[index][1] + h * MapSize[index][1] + w]);
          cls_index = 0;
        } else {
          for (int cl = 0; cl < class_num; cl++) {
            cls_val = cls[cl * MapSize[index][0] * MapSize[index][1] + h * MapSize[index][1] + w];

            if (0 == cl) {
              cls_max = cls_val;
              cls_index = cl;
            } else {
              if (cls_val > cls_max) {
                cls_max = cls_val;
                cls_index = cl;
              }
            }
          }
          cls_max = sigmoid(cls_max);
        }
        // printf("[liufeng debug] cls_max:%.4f\n",cls_max);
        if (cls_max < mdl_hb->cfg->draw_threshold) {continue;}

        RegDfl.clear();
        for (int lc = 0; lc < 4; lc++) {
          sfsum = 0;
          locval = 0;
          for (int df = 0; df < 16; df++) {
            locvaltemp = exp(reg[((lc * 16) + df) * MapSize[index][0] * MapSize[index][1] + h * MapSize[index][1] + w]);
            RegDeq[df] = locvaltemp;
            sfsum += locvaltemp;
          }
          for (int df = 0; df < 16; df++) {
            locvaltemp = RegDeq[df] / sfsum;
            locval += locvaltemp * df;
          }

          RegDfl.push_back(locval);
        }

        xmin = (MeshGrid[gridIndex + 0] - RegDfl[0]) * Strides[index] / w_scale;
        ymin = (MeshGrid[gridIndex + 1] - RegDfl[1]) * Strides[index] / h_scale;
        xmax = (MeshGrid[gridIndex + 0] + RegDfl[2]) * Strides[index] / w_scale;
        ymax = (MeshGrid[gridIndex + 1] + RegDfl[3]) * Strides[index] / h_scale;
        // std::cout << "xmin: " << xmin << " ymin: " << ymin << " xmax: " << xmax << " ymax: " << ymax << std::endl;
        // printf("[liufeng debug] %d %d\n",img_blob->ori_im_shape[1],img_blob->ori_im_shape[0]);
        xmin = xmin > 0 ? xmin : 0;
        ymin = ymin > 0 ? ymin : 0;
        xmax = xmax < img_blob->ori_im_shape[1] ? xmax : img_blob->ori_im_shape[1];
        ymax = ymax < img_blob->ori_im_shape[0] ? ymax : img_blob->ori_im_shape[0];
        

        labels.push_back(cls_index);
        confidences.push_back(cls_max);
        boxs.push_back(cv::Rect(xmin, ymin, xmax-xmin, ymax-ymin));
      }
    }
  }

  std::vector<int> indices;
  if (mdl_hb->cfg->arch == "YOLOV8") {
    cv::dnn::NMSBoxes(boxs, confidences, mdl_hb->cfg->draw_threshold, mdl_hb->cfg->nms_threshold, indices, 100);
  }
  // printf("[liufeng debug] indices.size:%d\n",indices.size());
  // V10 不需要nms
  if (mdl_hb->cfg->arch == "YOLOV10") {
    for (int idx=0; idx < boxs.size(); idx++) {indices.push_back(idx);}
  }

  // cv::Mat img;
  // nv122bgr(img_blob.yuv, cv::Size(1920,1080), img);  

  for (auto& idx : indices) {
    DetInfo* di = new DetInfo(confidences[idx], labels[idx], boxs[idx], mdl_hb->cfg->label_list[labels[idx]]);
    det_infos.emplace_back(di);
    // cv::rectangle(img, di->det_box, cv::Scalar(0, 0, 255), 2);
  }

  // cv::imwrite("res.jpg",img);

}
void HBDetPredictor::ParseTensor(float* data, int layer, 
  std::vector<cv::Rect>& boxs, std::vector<float>& scores, std::vector<int>& labels,
  std::vector<int> ori_shape, std::vector<int> feteature_shape, ModelConfig* cfg) {
    int num_classes = cfg->label_list.size();
    int stride = cfg->strides[layer];
    int num_pred = num_classes + 4 + 1;
    if (cfg->anchors_table.empty()) {
      spdlog::get("logger")->info("ERROR. HBDetPredictor::ParseTensor anchors_table is empty.");
      return;
    }
    std::vector<std::vector<int>> &anchors = cfg->anchors_table[layer];
    float w_scale = ori_shape[1] * 1.0 / cfg->input_shape[1];
    float h_scale = ori_shape[0] * 1.0 / cfg->input_shape[0];
    // spdlog::get("logger")->info("num_classes: {}, stride: {}, num_pred: {}, w_scale: {}, h_scale: {}", num_classes, stride, num_pred, w_scale, h_scale);

    int anchor_num = anchors.size();
    for (int h = 0; h < feteature_shape[0]; h++) {
      for (int w = 0; w < feteature_shape[1]; w++) {
        for (int k = 0; k < anchor_num; k++) {
          float *cur_data = data + k * num_pred;
          // if (h == 0 && w == 0 ) {printf("%f, %f, %f, %f, %f, %f, %f, %f\n", cur_data[0],cur_data[1],cur_data[2],cur_data[3],cur_data[4],cur_data[5],cur_data[6],cur_data[7]);}
          
          int id = std::distance(cur_data + 5, std::max_element(cur_data + 5, cur_data + 5 + num_classes));
          double x1 = 1 / (1 + std::exp(-cur_data[4])) * 1;
          double x2 = 1 / (1 + std::exp(-cur_data[id + 5]));
          double confidence = x1 * x2;
          // if (1 ) {printf("id:%d, 4:%f, id+5:%f, x1:%f, x2:%f, confidence:%f\n", id, cur_data[4], cur_data[id + 5], x1, x2, confidence);}
          if (confidence < cfg->draw_threshold) {continue;}
          // spdlog::get("logger")->info("confidence:{}", confidence);

          double cx = ((1.0 / (1.0 + std::exp(-cur_data[0]))) * 2 - 0.5 + w) * stride;
          double cy = ((1.0 / (1.0 + std::exp(-cur_data[1]))) * 2 - 0.5 + h) * stride;
          double ww = std::pow((1.0 / (1.0 + std::exp(-cur_data[2]))) * 2, 2) * anchors[k][0];
          double hh = std::pow((1.0 / (1.0 + std::exp(-cur_data[3]))) * 2, 2) * anchors[k][1];

          double xmin = (cx - ww / 2.0)* w_scale;
          double ymin = (cy - hh / 2.0)* h_scale;
          double xmax = (cx + ww / 2.0)* w_scale;
          double ymax = (cy + hh / 2.0)* h_scale;

          if (xmax <= 0 || ymax <= 0) { continue; }
          if (xmin > xmax || ymin > ymax) { continue; }
          if (xmin < 0) { xmin = 0;}
          if (ymin < 0) { ymin = 0;}
          if (xmax >= ori_shape[1]) { xmax = ori_shape[1]-1;}
          if (ymax >= ori_shape[0]) { ymax = ori_shape[0]-1;}

          boxs.emplace_back(cv::Rect(xmin, ymin, xmax-xmin, ymax-ymin));
          scores.emplace_back((float)confidence);
          labels.emplace_back(id);
        }
        data = data + num_pred * anchors.size();
      }
    }
}


int HBDetPredictor::get_tensor_hw(hbDNNTensor* tensor, int *height, int *width) {
  int h_index = 0;
  int w_index = 0;
  if (tensor->properties.tensorLayout == HB_DNN_LAYOUT_NHWC) {
    h_index = 1;
    w_index = 2;
  }
  else if (tensor->properties.tensorLayout == HB_DNN_LAYOUT_NCHW) {
    h_index = 2;
    w_index = 3;
  }
  else { return -1; }
  *height = tensor->properties.validShape.dimensionSize[h_index];
  *width = tensor->properties.validShape.dimensionSize[w_index];
  return 0;
}


}  // namespace hb
